// Created by zph on 2021/1/12.
// Modified by wzy on 2022/12/28.
// Modified by mzy on 2023/5/4.

#include <ros/ros.h>
#include <geometry_msgs/PoseStamped.h>
#include <mavros_msgs/State.h>
#include <string>
#include "math.h"
#include <offboard/ControlCommand.h>
#include <iostream>

#include "controller/colors.h"

/* -------------------------------------------------------------------------- */
/*                               Global Variable                              */
/* -------------------------------------------------------------------------- */
mavros_msgs::State cur_state;
geometry_msgs::PoseStamped cur_pose;
geometry_msgs::PoseStamped takeoff_pose;
geometry_msgs::PoseStamped hover_pose; // pose to hover in offboard mode

double hover_height = 1.0;
double r_circle = 1.0;
double t_total = 30.0;
double omege = 2 * M_PI / 30.0;
int count_circle = 0;
int count_total = 25 * 30;

int ctrl_rate = 25;

/* -------------------------------------------------------------------------- */
/*                     Function Declaration and Definition                    */
/* -------------------------------------------------------------------------- */

// Tool Function
template <typename T>
void readParam(ros::NodeHandle &nh, std::string param_name, T &loaded_param)
{
  // template to read param from roslaunch
  const std::string &node_name = ros::this_node::getName();
  param_name = node_name + "/" + param_name;
  if (!nh.getParam(param_name, loaded_param))
  {
    std::cout << RED_B << "Fail to load " << param_name
              << ", use default value:" << loaded_param
              << RESET << std::endl;
  }
  else
  {
    std::cout << GREEN_B << "Load successfully " << param_name
              << ": " << loaded_param
              << RESET << std::endl;
  }
}

offboard::ControlCommand pose2pva(geometry_msgs::PoseStamped pose)
{
  offboard::ControlCommand pva;
  pva.header.stamp = ros::Time::now();
  pva.pos.x = pose.pose.position.x;
  pva.pos.y = pose.pose.position.y;
  pva.pos.z = pose.pose.position.z;
  pva.vel.x = 0.0;
  pva.vel.y = 0.0;
  pva.vel.z = 0.0;
  pva.acc.x = 0.0;
  pva.acc.y = 0.0;
  pva.acc.z = 0.0;
  pva.yaw = 0.0;
  return pva;
}

void loadRosParams(ros::NodeHandle &nh)
{
  readParam<int>(nh, "ctrl_rate", ctrl_rate);
  readParam<double>(nh, "hover_height", hover_height);
  readParam<double>(nh, "r_circle", r_circle);
  readParam<double>(nh, "t_total", t_total);
}

// callback function
void uav_state_cb(const mavros_msgs::State::ConstPtr &msg)
{
  cur_state = *msg;
}

void uav_pose_cb(const geometry_msgs::PoseStamped::ConstPtr &msg)
{
  cur_pose = *msg;
}

/* -------------------------------------------------------------------------- */
/*                                main function                               */
/* -------------------------------------------------------------------------- */

int main(int argc, char **argv)
{
  // ros init
  ros::init(argc, argv, "offboard_node", ros::init_options::AnonymousName);
  ros::NodeHandle nh;

  // load param
  loadRosParams(nh);

  // circle param
  omege = 2 * M_PI / t_total;
  count_circle = 0;
  count_total = int(ctrl_rate * t_total);

  // ros pub and sub
  ros::Subscriber uav_state_sub = nh.subscribe<mavros_msgs::State>(
      "/mavros/state", 10, uav_state_cb);
  ros::Subscriber uav_pose_sub = nh.subscribe<geometry_msgs::PoseStamped>(
      "/mavros/vision_pose/pose", 2, uav_pose_cb);
  ros::Publisher setpva_pub = nh.advertise<offboard::ControlCommand>(
      "/setpoint_pva", 10);

  // the setpoint publishing rate MUST be faster than 2Hz
  ros::Rate ctrl_loop(ctrl_rate);

  // wait for FCU connection
  while (ros::ok() && !cur_state.connected)
  {
    ros::spinOnce();
    ROS_INFO_THROTTLE(1, "wait for FCU connection");
    ctrl_loop.sleep();
  }

  /**
   * Pre Step :
   * Send a few setpoints before starting
   * Before entering Offboard mode,
   * you must have already started streaming setpoints.
   * Otherwise the mode switch will be rejected.
   */
  for (int i = 10; ros::ok() && i > 0; --i)
  {
    ros::spinOnce();
    setpva_pub.publish(pose2pva(cur_pose));
    ctrl_loop.sleep();
  }

  // takeoff: Record whether the uav has taken off
  // just_takeoff: check whether the uav change into offboard mode from other mode
  bool takeoff = false;
  bool just_takeoff = true;

  /**************************** main ctrl loop *******************************/
  while (ros::ok())
  {
    ctrl_loop.sleep();
    ros::spinOnce();

    if (!cur_state.armed) {
      ROS_INFO_STREAM_THROTTLE(1, " waiting for Vehicle arm");
      setpva_pub.publish(pose2pva(cur_pose));
      takeoff = false;
      just_takeoff = true;
      continue;
    }

    if (cur_state.mode != "OFFBOARD")
    {
      ROS_INFO_STREAM_THROTTLE(1, " waiting for Offboard cmd");
      setpva_pub.publish(pose2pva(cur_pose));
      just_takeoff = true;
      continue;
    }

    ROS_INFO_STREAM_THROTTLE(1, "\033[33m OFFBOARD Mode \033[0m");

    // Take off to specific height
    if (!takeoff)
    {
      takeoff_pose.pose.position = cur_pose.pose.position;
      takeoff_pose.pose.position.z = hover_height;

      geometry_msgs::PoseStamped climb_pose;
      climb_pose.pose.position = cur_pose.pose.position;

      double err = 0.1; // TODO: 高度阈值
      while (ros::ok() &&
             fabs(takeoff_pose.pose.position.z - cur_pose.pose.position.z) > err)
      {
        ros::spinOnce();
        // TODO:上升速度 0.2m/s
        climb_pose.pose.position.z = climb_pose.pose.position.z +
                                     0.2 / double(ctrl_rate);
        climb_pose.pose.position.z = std::min(climb_pose.pose.position.z,
                                              takeoff_pose.pose.position.z);

        setpva_pub.publish(pose2pva(climb_pose));
        ROS_INFO("[Takeoff]: Target: %.3f, %.3f, %.3f | Curr: %.3f, %.3f, %.3f",
                 climb_pose.pose.position.x,
                 climb_pose.pose.position.y,
                 climb_pose.pose.position.z,
                 cur_pose.pose.position.x,
                 cur_pose.pose.position.y,
                 cur_pose.pose.position.z);
        ctrl_loop.sleep();
      }
      ROS_INFO_STREAM("\033[33m Take off done \033[0m");
      takeoff = true;
      hover_pose = takeoff_pose;
    }

    // Switch into offboard mode just now, stay still for a while
    if (just_takeoff)
    {
      ROS_INFO_STREAM("\033[33m Hovering for several seconds \033[0m");
      for (int i = ctrl_rate * 4; ros::ok() && i > 0; --i)
      {
        ros::spinOnce();
        setpva_pub.publish(pose2pva(takeoff_pose));
        ctrl_loop.sleep();
      }
      just_takeoff = false;
      hover_pose = takeoff_pose;
    }

    /**
     * x = r*sin(w*t)
     * y = r*(1-cos(w*t))
     * vx = r*w*cos(w*t)
     * vy = r*w*sin(w*t)
     * ax = -r*w^2*sin(w*t)
     * ay = r*w^2*cos(w*t)
     *
     * t = n/k, k为控制频率，n为计数点
     */
    ROS_INFO_STREAM_ONCE("\033[33m In Offboard mode, start circling! \033[0m");
    if (count_circle < count_total + 1)
    {
      offboard::ControlCommand pva_traj;
      pva_traj.header.frame_id = "global";
      pva_traj.header.stamp = ros::Time::now();
      double phi = double(count_circle) / double(ctrl_rate) * omege;

      /**
       *  Here is the code for testing.
       *  When running a specific case,
       *  please comment out the other cases.
       */

      /**
       *  case 0 hovering
       */
      pva_traj.pos.x = 0 + takeoff_pose.pose.position.x;
      pva_traj.pos.y = 0 + takeoff_pose.pose.position.y;
      pva_traj.pos.z = hover_height;
      pva_traj.vel.x = 0;
      pva_traj.vel.y = 0;
      pva_traj.vel.z = 0;
      pva_traj.acc.x = 0;
      pva_traj.acc.y = 0;
      pva_traj.acc.z = 0;
      pva_traj.yaw = 0.0;

      /**
       *  case 1 circling WITHOUT yaw rotating
       */
      // pva_traj.pos.x = r_circle * sin(phi) +
      //                  takeoff_pose.pose.position.x;
      // pva_traj.pos.y = r_circle * (1 - cos(phi)) +
      //                  takeoff_pose.pose.position.y;
      // pva_traj.pos.z = hover_height;
      // pva_traj.vel.x = 0;
      // pva_traj.vel.y = 0;
      // pva_traj.vel.z = 0;
      // pva_traj.acc.x = 0;
      // pva_traj.acc.y = 0;
      // pva_traj.acc.z = 0;
      // pva_traj.yaw = 0.0;

      /**
       *  case 2  yaw rotating WITHOUT moving
       */
      // pva_traj.pos.x = 0 + takeoff_pose.pose.position.x;
      // pva_traj.pos.y = 0 + takeoff_pose.pose.position.y;
      // pva_traj.pos.z = hover_height;
      // pva_traj.vel.x = 0;
      // pva_traj.vel.y = 0;
      // pva_traj.vel.z = 0;
      // pva_traj.acc.x = 0;
      // pva_traj.acc.y = 0;
      // pva_traj.acc.z = 0;
      // pva_traj.yaw = phi;

      /**
       *  case 3 circling WITH yaw rotating
       */
      // pva_traj.pos.x = r_circle * sin(phi) + takeoff_pose.pose.position.x;
      // pva_traj.pos.y = r_circle * (1 - cos(phi)) + takeoff_pose.pose.position.y;
      // pva_traj.pos.z = hover_height;
      // pva_traj.vel.x = 0;
      // pva_traj.vel.y = 0;
      // pva_traj.vel.z = 0;
      // pva_traj.acc.x = 0;
      // pva_traj.acc.y = 0;
      // pva_traj.acc.z = 0;
      // pva_traj.yaw = phi;

      /**
       *  case 4 circling WITHOUT yaw rotating,
       *         specify velocity and accleration
       */
      // pva_traj.pos.x = r_circle * sin(phi) + takeoff_pose.pose.position.x;
      // pva_traj.pos.y = r_circle * (1 - cos(phi)) + takeoff_pose.pose.position.y;
      // pva_traj.pos.z = hover_height;
      // pva_traj.vel.x = r_circle * omege * cos(phi);
      // pva_traj.vel.y = r_circle * omege * sin(phi);
      // pva_traj.vel.z = 0;
      // pva_traj.acc.x = -r_circle * pow(omege, 2) * sin(phi);
      // pva_traj.acc.y = r_circle * pow(omege, 2) * cos(phi);
      // pva_traj.acc.z = 0;
      // pva_traj.yaw = 0.0;

      /**
       *  case 5 circling WITH yaw rotating,
       *         specify velocity and accleration
       */
      // pva_traj.pos.x = r_circle * sin(phi) + takeoff_pose.pose.position.x;
      // pva_traj.pos.y = r_circle * (1 - cos(phi)) + takeoff_pose.pose.position.y;
      // pva_traj.pos.z = hover_height;
      // pva_traj.vel.x = r_circle * omege * cos(phi);
      // pva_traj.vel.y = r_circle * omege * sin(phi);
      // pva_traj.vel.z = 0;
      // pva_traj.acc.x = -r_circle * pow(omege, 2) * sin(phi);
      // pva_traj.acc.y = r_circle * pow(omege, 2) * cos(phi);
      // pva_traj.acc.z = 0;
      // pva_traj.yaw = phi;

      setpva_pub.publish(pva_traj);

      hover_pose = cur_pose;
      count_circle++;
      ROS_INFO("[traj] Target: %.3f, %.3f, %.3f | Curr: %.3f, %.3f, %.3f",
               pva_traj.pos.x,
               pva_traj.pos.y,
               pva_traj.pos.z,
               cur_pose.pose.position.x,
               cur_pose.pose.position.y,
               cur_pose.pose.position.z);
    }
    else
    {

      ROS_WARN_THROTTLE(1, " Trajectory finished, Hover ... ");
      setpva_pub.publish(pose2pva(hover_pose));
    }
  }

  return 0;
}
